package redshift

import (
	"log/slog"

	"github.com/antlr4-go/antlr/v4"
	"github.com/pkg/errors"

	parser "github.com/bytebase/parser/redshift"

	"github.com/bytebase/bytebase/backend/common"
	storepb "github.com/bytebase/bytebase/backend/generated-go/store"
	"github.com/bytebase/bytebase/backend/plugin/parser/base"
)

func init() {
	// Unregister the PostgreSQL splitter for Redshift
	// and register the Redshift-specific splitter
	base.RegisterSplitterFunc(storepb.Engine_REDSHIFT, SplitSQL)
}

// SplitSQL splits the given SQL statement into multiple SQL statements.
func SplitSQL(statement string) ([]base.SingleSQL, error) {
	lexer := parser.NewRedshiftLexer(antlr.NewInputStream(statement))
	stream := antlr.NewCommonTokenStream(lexer, antlr.TokenDefaultChannel)
	list, err := splitSQLImpl(stream, statement)
	if err != nil {
		slog.Info("failed to split Redshift statement", "statement", statement)
		// Use parser to split statement.
		return splitByParser(statement, lexer, stream)
	}
	return list, nil
}

func splitByParser(statement string, lexer *parser.RedshiftLexer, stream *antlr.CommonTokenStream) ([]base.SingleSQL, error) {
	p := parser.NewRedshiftParser(stream)
	lexerErrorListener := &base.ParseErrorListener{
		Statement: statement,
	}
	lexer.RemoveErrorListeners()
	lexer.AddErrorListener(lexerErrorListener)

	parserErrorListener := &base.ParseErrorListener{
		Statement: statement,
	}
	p.RemoveErrorListeners()
	p.AddErrorListener(parserErrorListener)

	p.BuildParseTrees = true

	tree := p.Root()
	if lexerErrorListener.Err != nil {
		return nil, lexerErrorListener.Err
	}

	if parserErrorListener.Err != nil {
		return nil, parserErrorListener.Err
	}

	if tree == nil || tree.Stmtblock() == nil || tree.Stmtblock().Stmtmulti() == nil {
		return nil, errors.New("failed to split multiple statements")
	}

	var result []base.SingleSQL
	tokens := stream.GetAllTokens()

	start := 0
	for _, semi := range tree.Stmtblock().Stmtmulti().AllSEMI() {
		pos := semi.GetSymbol().GetStart()
		antlrPosition := base.FirstDefaultChannelTokenPosition(tokens[start : pos+1])
		// From antlr4, the line is ONE based, and the column is ZERO based.
		// So we should minus 1 for the line.
		result = append(result, base.SingleSQL{
			Text:     stream.GetTextFromTokens(tokens[start], tokens[pos]),
			BaseLine: tokens[start].GetLine() - 1,
			End: common.ConvertANTLRPositionToPosition(&common.ANTLRPosition{
				Line:   int32(tokens[pos].GetLine()),
				Column: int32(tokens[pos].GetColumn()),
			}, statement),
			Start: common.ConvertANTLRPositionToPosition(antlrPosition, statement),
			Empty: base.IsEmpty(tokens[start:pos+1], parser.RedshiftParserSEMI),
		})
		start = pos + 1
	}
	// For the last statement, it may not end with semicolon symbol, EOF symbol instead.
	eofPos := len(tokens) - 1
	if start < eofPos {
		antlrPosition := base.FirstDefaultChannelTokenPosition(tokens[start:])
		// From antlr4, the line is ONE based, and the column is ZERO based.
		// So we should minus 1 for the line.
		result = append(result, base.SingleSQL{
			Text:     stream.GetTextFromTokens(tokens[start], tokens[eofPos-1]),
			BaseLine: tokens[start].GetLine() - 1,
			End: common.ConvertANTLRPositionToPosition(&common.ANTLRPosition{
				Line:   int32(tokens[eofPos-1].GetLine()),
				Column: int32(tokens[eofPos-1].GetColumn()),
			}, statement),
			Start: common.ConvertANTLRPositionToPosition(antlrPosition, statement),
			Empty: base.IsEmpty(tokens[start:eofPos], parser.RedshiftParserSEMI),
		})
	}
	return result, nil
}

type openParenthesis struct {
	tokenType int
	pos       int
}

func splitSQLImpl(stream *antlr.CommonTokenStream, statement string) ([]base.SingleSQL, error) {
	var result []base.SingleSQL
	stream.Fill()
	tokens := stream.GetAllTokens()

	var beginCaseStack, ifStack, loopStack []openParenthesis
	var semicolonStack []int

	for i, token := range tokens {
		switch token.GetTokenType() {
		case parser.RedshiftParserBEGIN_P:
			if isBeginTransaction(tokens, i) {
				continue
			}

			beginCaseStack = append(beginCaseStack, openParenthesis{
				tokenType: token.GetTokenType(),
				pos:       i,
			})
		case parser.RedshiftParserCASE:
			beginCaseStack = append(beginCaseStack, openParenthesis{
				tokenType: token.GetTokenType(),
				pos:       i,
			})
		case parser.RedshiftParserIF_P:
			ifStack = append(ifStack, openParenthesis{
				tokenType: token.GetTokenType(),
				pos:       i,
			})
		case parser.RedshiftParserLOOP:
			loopStack = append(loopStack, openParenthesis{
				tokenType: token.GetTokenType(),
				pos:       i,
			})
		case parser.RedshiftParserSEMI:
			semicolonStack = append(semicolonStack, i)
		case parser.RedshiftParserEND_P:
			if isEndTransaction(tokens, i) {
				continue
			}

			nextToken := base.GetDefaultChannelTokenType(tokens, i, 1)
			switch nextToken {
			case parser.RedshiftParserIF_P:
				if len(ifStack) == 0 {
					return nil, errors.New("invalid statement: failed to split multiple statements")
				}
				// There are two cases:
				// 1. The IF statement with END IF statement.
				// 2. The IF statement without END IF statement.
				// We should match the longest IF statement, so we should check the first IF statement in the stack.
				semicolonStack = popSemicolonStack(semicolonStack, ifStack[0].pos)
				ifStack = ifStack[:len(ifStack)-1]
			case parser.RedshiftParserLOOP:
				if len(loopStack) == 0 {
					return nil, errors.New("invalid statement: failed to split multiple statements")
				}
				semicolonStack = popSemicolonStack(semicolonStack, loopStack[len(semicolonStack)-1].pos)
				loopStack = loopStack[:len(loopStack)-1]
			case parser.RedshiftParserCASE:
				if len(beginCaseStack) == 0 {
					return nil, errors.New("invalid statement: failed to split multiple statements")
				}
				semicolonStack = popSemicolonStack(semicolonStack, beginCaseStack[len(beginCaseStack)-1].pos)
				beginCaseStack = beginCaseStack[:len(beginCaseStack)-1]
			default:
				if len(beginCaseStack) == 0 {
					return nil, errors.New("invalid statement: failed to split multiple statements")
				}
				semicolonStack = popSemicolonStack(semicolonStack, beginCaseStack[len(beginCaseStack)-1].pos)
				beginCaseStack = beginCaseStack[:len(beginCaseStack)-1]
			}
		default:
		}
	}

	start := 0
	for _, pos := range semicolonStack {
		antlrPosition := base.FirstDefaultChannelTokenPosition(tokens[start : pos+1])
		// From antlr4, the line is ONE based, and the column is ZERO based.
		// So we should minus 1 for the line.
		result = append(result, base.SingleSQL{
			Text:     stream.GetTextFromTokens(tokens[start], tokens[pos]),
			BaseLine: tokens[start].GetLine() - 1,
			End: common.ConvertANTLRPositionToPosition(&common.ANTLRPosition{
				Line:   int32(tokens[pos].GetLine()),
				Column: int32(tokens[pos].GetColumn()),
			}, statement),
			Start: common.ConvertANTLRPositionToPosition(antlrPosition, statement),
			Empty: base.IsEmpty(tokens[start:pos+1], parser.RedshiftParserSEMI),
		})
		start = pos + 1
	}
	// For the last statement, it may not end with semicolon symbol, EOF symbol instead.
	eofPos := len(tokens) - 1
	if start < eofPos {
		antlrPosition := base.FirstDefaultChannelTokenPosition(tokens[start:])
		// From antlr4, the line is ONE based, and the column is ZERO based.
		// So we should minus 1 for the line.
		result = append(result, base.SingleSQL{
			Text:     stream.GetTextFromTokens(tokens[start], tokens[eofPos-1]),
			BaseLine: tokens[start].GetLine() - 1,
			End: common.ConvertANTLRPositionToPosition(&common.ANTLRPosition{
				Line:   int32(tokens[eofPos-1].GetLine()),
				Column: int32(tokens[eofPos-1].GetColumn()),
			}, statement),
			Start: common.ConvertANTLRPositionToPosition(antlrPosition, statement),
			Empty: base.IsEmpty(tokens[start:eofPos], parser.RedshiftParserSEMI),
		})
	}

	return result, nil
}

func popSemicolonStack(semicolonStack []int, pos int) []int {
	if len(semicolonStack) == 0 {
		return semicolonStack
	}

	for i := len(semicolonStack) - 1; i >= 0; i-- {
		if semicolonStack[i] < pos {
			return semicolonStack[:i+1]
		}
	}

	return []int{}
}

func isEndTransaction(tokens []antlr.Token, index int) bool {
	if tokens[index].GetTokenType() != parser.RedshiftLexerEND_P {
		return false
	}

	switch base.GetDefaultChannelTokenType(tokens, index, 1) {
	case parser.RedshiftParserTRANSACTION,
		parser.RedshiftParserWORK,
		parser.RedshiftParserSEMI:
		return true
	case parser.RedshiftParserAND:
		if base.GetDefaultChannelTokenType(tokens, index, 2) == parser.RedshiftParserNO {
			return base.GetDefaultChannelTokenType(tokens, index, 3) == parser.RedshiftParserCHAIN
		}
		return base.GetDefaultChannelTokenType(tokens, index, 2) == parser.RedshiftParserCHAIN
	default:
		return false
	}
}

func isBeginTransaction(tokens []antlr.Token, index int) bool {
	if tokens[index].GetTokenType() != parser.RedshiftLexerBEGIN_P {
		return false
	}

	switch base.GetDefaultChannelTokenType(tokens, index, 1) {
	case parser.RedshiftParserTRANSACTION,
		parser.RedshiftParserWORK,
		parser.RedshiftParserSEMI,
		parser.RedshiftParserISOLATION,
		parser.RedshiftParserREAD,
		parser.RedshiftLexerDEFERRABLE:
		return true
	case parser.RedshiftParserNOT:
		return base.GetDefaultChannelTokenType(tokens, index, 2) == parser.RedshiftParserDEFERRABLE
	default:
		return false
	}
}
